import java.util.Arrays;
import java.util.Comparator;
import java.util.HashMap;
import java.util.Map;
import java.util.PriorityQueue;

/**
 * 347. 前 K 个高频元素
 */
public class Solution_347 {
    /**
     * 堆
     * <p>
     * 时间复杂度：O(Nlogk)
     * <p>
     * 空间复杂度：O(N)
     */
    public int[] topKFrequent(int[] nums, int k) {
        Map<Integer, Integer> occurrences = new HashMap<>();
        for (int num : nums) {
            occurrences.put(num, occurrences.getOrDefault(num, 0) + 1);
        }

        // 小顶堆
        // int[] 的第一个元素代表数组的值，第二个元素代表了该值出现的次数
        PriorityQueue<int[]> heap = new PriorityQueue<int[]>(k, new Comparator<int[]>() {
            @Override
            public int compare(int[] o1, int[] o2) {
                return o1[1] - o2[1];
            }
        });
        for (Map.Entry<Integer, Integer> entry : occurrences.entrySet()) {
            int num = entry.getKey(), count = entry.getValue();
            if (heap.size() == k) {
                // 如果堆的元素个数等于 k，则检查堆顶与当前出现次数的大小
                if (heap.peek()[1] < count) {
                    heap.poll();
                    heap.offer(new int[] { num, count });
                }
            } else {
                // 如果堆的元素个数小于 k，就可以直接插入堆中
                heap.offer(new int[] { num, count });
            }
        }

        int[] ans = new int[k];
        for (int i = 0; i < k; i++) {
            ans[i] = heap.poll()[0];
        }
        return ans;
    }

    public static void main(String[] args) {
        Solution_347 solution = new Solution_347();
        int[] nums = { 1, 1, 1, 2, 2, 3 };
        int k = 2;
        int[] ans = solution.topKFrequent(nums, k);
        System.out.println(Arrays.toString(ans));
    }
}
